/* Copyright (c) 2024 Huawei Technologies Co., Ltd.
 * This file is a part of the CANN Open Software.
 * Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
 * Please refer to the License for details. You may not use this file except in compliance with the License.
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
 * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
 * See LICENSE in the root of the software repository for the full text of the License.
 * ===================================================================================================================*/

#ifndef GRAPH_ASCEND_IR_H
#define GRAPH_ASCEND_IR_H

#include <string>
#include <memory>
#include "attr_store.h"
#include "graph/compute_graph.h"
#include "graph/ascend_ir/expression.h"
#include "graph/node.h"
#include "graph/anchor.h"
#include "debug/ge_util.h"
#include "graph/utils/op_desc_utils.h"
#include "common/ge_common/debug/ge_log.h"
#include "external/graph/operator.h"

namespace ge {
namespace {
constexpr int64_t kIdNone = -1;
}

class AscIRException : public std::exception {
 public:
  struct Info {
    graphStatus error_code;
  };
  AscIRException(const Info &info);
  const Info &GetInfo() const;
 private:
   Info info_;
};

#define CHECK_NOTNULL_WITH_THROW_EXCEPTION(val, ...)                                        \
  do {                                                                                      \
    if ((val) == nullptr) {                                                                 \
      GELOGE(ge::GRAPH_FAILED, "[Check][Param:" #val "]null is invalid" __VA_ARGS__);       \
      throw AscIRException({ge::GRAPH_FAILED});                                             \
    }                                                                                       \
  } while (false)

#define CHECK_BOOL_WITH_THROW_EXCEPTION(error_core, val, ...)                               \
  do {                                                                                      \
    if (!(val)) {                                                                           \
      GELOGE(error_core, "[Check][Expr:" #val "] is false" __VA_ARGS__);                    \
      throw AscIRException({error_core});                                                   \
    }                                                                                       \
  } while (false)

struct Axis {
  using Type = enum : int32_t {
    AXIS_TYPE_ORIGINAL,
    AXIS_TYPE_BLOCK_OUTER,  // outer axis after split by multicore
    AXIS_TYPE_BLOCK_INNER,  // inner axis after split by multicore
    AXIS_TYPE_TILE_OUTER,   // outer axis after split by one core
    AXIS_TYPE_TILE_INNER,   // inner axis after split by one core
    AXIS_TYPE_MERGED
  };

  int64_t id;        // axis id
  std::string name;  // axis name
  Type type;
  bool bind_block;
  ge::expression::ExpressionPtr size;
  int32_t align;
  std::vector<int64_t> from;
  uint64_t split_pair_other_id;
  bool allow_oversize_axis;
  bool allow_unaligned_tail;
  std::pair<int64_t, int64_t> value_range = {-1, -1};
};
using AxisPtr = std::shared_ptr<Axis>;
using AxisId = int64_t;

struct Optional {
  uint32_t idx;
  std::string name;
  std::string data_type;
  std::string min_value;
  std::string max_value;
};
using OptionalPtr = std::shared_ptr<Optional>;

enum class ComputeType : int32_t {
  COMPUTE_DATA,
  COMPUTE_REDUCE_DATA,
  COMPUTE_WORKSPACE,
  COMPUTE_LOAD,
  COMPUTE_STORE,
  COMPUTE_REDUCE_STORE,
  COMPUTE_ELEWISE,
  COMPUTE_BROADCAST,
  COMPUTE_REDUCE,
  COMPUTE_TRANPOSE
};

enum class ComputeUnit : int32_t {
  UNIT_NONE,
  UNIT_MTE1,
  UNIT_MTE2,
  UNIT_MTE3,
  UNIT_SCALAR,
  UNIT_VECTOR,
  UNIT_CUBE,
};

enum class ApiType : int32_t {
  API_TYPE_BUFFER,
  API_TYPE_COMPUTE
};

struct ApiInfo {
  ApiType type;
  ComputeUnit unit;
};

struct HintComputeType {
  ComputeType compute_type;
};

struct SchedInfo {
  int64_t exec_order;
  int64_t sub_exec_order;
  std::vector<int64_t> axis;
  int64_t loop_axis;
  std::vector<int64_t> initial_axis;
};

struct AscNodeAttr : public ge::AttrGroupsBase {
  SchedInfo sched;
  ApiInfo api;
  HintComputeType hint;
  AscNodeAttr() = default;
  virtual ~AscNodeAttr() = default;
  static AscNodeAttr &GetAttrGroup(ge::Operator& op);
  AttrGroupsBase *CreateNew() override;
};

enum class AllocType : int32_t { ALLOC_TYPE_GLOBAL, ALLOC_TYPE_L1, ALLOC_TYPE_L2, ALLOC_TYPE_BUFFER, ALLOC_TYPE_QUEUE };

enum class MemHardware : int32_t {
  MEM_HARDWARE_GM,
  MEM_HARDWARE_UB,
};

enum class Position : int32_t {
  POSITION_GM,
  POSITION_VECIN,
  POSITION_VECOUT,
};

struct MemAttr {
  int64_t tensor_id = -1;
  AllocType alloc_type = AllocType::ALLOC_TYPE_GLOBAL;
  Position position = Position::POSITION_GM;
  MemHardware hardware = MemHardware::MEM_HARDWARE_GM;
  std::vector<int64_t> buf_ids;
  std::string name;
};

struct MemQueueAttr {
  int64_t id;
  int64_t depth;
  int64_t buf_num;
  std::string name;
};

struct MemBufAttr {
  int64_t id;
  std::string name;
};

struct MemOptAttr {
  int64_t reuse_id;
  int64_t ref_tensor;
  int64_t merge_scope;
};

struct AscNodeOutputs;
class AscTensorDataType {
 public:
  operator ge::DataType() const {
    CHECK_NOTNULL_WITH_THROW_EXCEPTION(tensor_desc_);
    return tensor_desc_->GetDataType();
  };
 private:
  friend struct AscNodeOutputs;
  GeTensorDesc *tensor_desc_;
};

class AscTensorFormat {
 public:
  operator ge::Format() const {
    CHECK_NOTNULL_WITH_THROW_EXCEPTION(tensor_desc_);
    return tensor_desc_->GetFormat();
  };
 private:
  friend struct AscNodeOutputs;
  GeTensorDesc *tensor_desc_;
};

struct AscTensorAttr : public ge::AttrGroupsBase {
  AscTensorDataType dtype;
  AscTensorFormat format;
  std::vector<int64_t> axis;
  std::vector<ge::expression::ExpressionPtr> repeats;
  std::vector<ge::expression::ExpressionPtr> strides;
  std::vector<int64_t> vectorized_axis;
  std::vector<ge::expression::ExpressionPtr> vectorized_strides;
  MemAttr mem;
  MemQueueAttr que;
  MemBufAttr buf;
  MemOptAttr opt;
  AttrGroupsBase *CreateNew() override;
  static AscTensorAttr &GetTensorAttr(ge::Operator *op, const uint32_t index);
};

struct AscNodeOutputs {
  explicit AscNodeOutputs(const ge::OpDescPtr &op_desc);
  AscTensorAttr &operator[](uint32_t index);
  std::vector<AscTensorAttr *> operator()();
 private:
  ge::OpDescPtr op_desc_;
};

struct AscNodeInputs {
  explicit AscNodeInputs(ge::Node *node);
  AscTensorAttr &operator[](uint32_t index);
  uint32_t Size();
 private:
  ge::Node *node_;
};

class AscNode : public Node {
public:
  AscNode(const OpDescPtr &op_desc, const ComputeGraphPtr &compute_graph);
  AscNodeInputs inputs;
  AscNodeOutputs outputs;
  AscNodeAttr &attr;
};
using AscNodePtr = std::shared_ptr<AscNode>;

class AscNodeIter {
 public:
  explicit AscNodeIter(ge::ComputeGraph::Vistor<ge::NodePtr>::Iterator &&iter);
  AscNodeIter &operator++();
  AscNodePtr operator*();
  bool operator!=(const AscNodeIter &other) const;
 private:
  ge::ComputeGraph::Vistor<ge::NodePtr>::Iterator impl_;
};

class AscNodeVisitor {
 public:
  using Iterator = ge::ComputeGraph::Vistor<ge::NodePtr>::Iterator;
  AscNodeIter begin();
  AscNodeIter end();
  explicit AscNodeVisitor(ge::ComputeGraph::Vistor<ge::NodePtr> &&visitor);
 private:
  ge::ComputeGraph::Vistor<ge::NodePtr> impl_;
};

class AscGraphImpl;
namespace ascir {
namespace cg {
class CodeGenUtils;
}
}
class AscGraph {
  friend class ascir::cg::CodeGenUtils;
 public:
  explicit AscGraph(const char *name);
  ~AscGraph();
  void SetTilingKey(const uint32_t tiling_key);
  int64_t GetTilingKey();
  Axis &CreateAxis(const std::string &name, const ge::expression::ExpressionPtr &size,
                   const int64_t min_value = -1, const int64_t max_value = -1);
  Axis *FindAxis(const int64_t axis_id);
  Optional &CreateOptionalAtt(const std::string &name, uint32_t idx, int32_t min_value, int32_t max_value);
  Optional &CreateOptionalAtt(const std::string &name, uint32_t idx, uint32_t min_value, uint32_t max_value);
  Optional &CreateOptionalAtt(const std::string &name, uint32_t idx, float min_value, float max_value);
  Optional &CreateOptionalAtt(const std::string &name, uint32_t idx, int8_t min_value, int8_t max_value);
  Optional &CreateOptionalAtt(const std::string &name, uint32_t idx, const std::string &data_type,
                              const std::string &min_value, const std::string &max_value);
  void AddNode(ge::Operator &op);
  AscNodePtr FindNode(const char *name) const;
  std::pair<AxisPtr, AxisPtr> BlockSplit(const int64_t axis_id, const std::string &inner_axis_name = "",
                                         const std::string &outer_axis_name = "");
  std::pair<AxisPtr, AxisPtr> TileSplit(const int64_t axis_id, const std::string &inner_axis_name = "",
                                        const std::string &outer_axis_name = "");
  AxisPtr MergeAxis(const std::vector<int64_t> &axis_ids, const std::string &merge_axis_name = "");
  void ApplySplit(const AscNodePtr &node, const int64_t outter_id, const int64_t inner_id);
  void ApplyMerge(const AscNodePtr &node, const int64_t merged_axis_id);
  void ApplyReorder(const AscNodePtr &node, const std::vector<int64_t> &reordered_axis);
  AscNodeVisitor GetAllNodes() const;
  AscNodeVisitor GetInputNodes() const;
  std::vector<AxisPtr> GetAllAxis() const;
  std::vector<OptionalPtr> GetAllOptional() const;
  std::string GetName() const;
 private:
  std::shared_ptr<AscGraphImpl> impl_;
};

void AddEdgeForNode(ge::Operator &src_op, int32_t src_index, ge::Operator &dst_op,
                    int32_t dst_index);
}  // namespace ge

#endif  // GRAPH_ASCEND_IR_H
